# 1.实现LeNet-5卷积神经网络
import torch
# 6.
import torch.nn.functional as F

# 8.
# 导入summary
# from torchsummary import summary

# nn.relu 和F.relu其实是一样的
# F.relu一般用于函数调用，使用在forward中，而nn.relu是模块调用，一般定义在网络层(__init__中)
# 如果使用summary 使用nn.relu会打印出，F.relu不会打印出

# device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# 定义神经网络类
class LeNet5(torch.nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 卷积
        # 函数原型解析 看pytorch.org文档
        # torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True,
        #  padding_mode='zeros', device=None, dtype=None)
        #
        # in_channels：输入图像的通道数。
        # out_channels：经过卷积运算产生的通道数。
        # kernel_size：卷积核大小，整数或元组类型。
        # stride：卷积运算的步幅，整数或元组类型，默认为1。
        # padding：边界填充值，整数或元组类型，四周进行填充，默认为0。
        # 为啥是padding为2是因为输入的图片是32x32的图片，Mnist数据集是28x28的，组成32x32的图片
        self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=6, kernel_size=(5, 5), stride=1, padding=2)
        # 池化函数看pytorch.org文档 有kernel_size=2 和步长stride=2
        self.pool1 = torch.nn.AvgPool2d(kernel_size=2, stride=2)
        # 再次池化
        self.conv2 = torch.nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        self.pool2 = torch.nn.AvgPool2d(kernel_size=2, stride=2)
        # 全连接
        self.fc1 = torch.nn.Linear(in_features=400, out_features=120)
        self.fc2 = torch.nn.Linear(in_features=120, out_features=84)
        self.fc3 = torch.nn.Linear(in_features=84, out_features=10)

    def forward(self, x):
        # 要经过第一次的卷积后跟激活函数这里使用relu
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        # 进行全连接前要先进行展平使用view进行展平,改变张量形状
        x = x.view(-1, 400)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        # pytorch交叉熵自带softmax如果你没使用softmax会自动使用softmax
        return x
